# TODO: Deprecate in favor of AtomWorks
# (This code is a holdover from previous frameworks)

import os

import torch

script_dir = os.path.dirname(os.path.realpath(__file__)) + "/"

NHEAVY = 23
NTOTAL = 36
NNAPROTAAS = 32
NAATOKENS = 20 + 2 + 10 + 1 + 47  # 20 AAs, UNK, MASK, 8 NAs,HIS_D, 47 atoms
costgtNA = -0.2744

aa2long = [
    (  # 0 ala
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "3HB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 0  ala
    (  # 1 arg
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD ",
        " NE ",
        " CZ ",
        " NH1",
        " NH2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HG ",
        "2HG ",
        "1HD ",
        "2HD ",
        " HE ",
        "1HH1",
        "2HH1",
        "1HH2",
        "2HH2",
    ),  # 1  arg
    (  # 2  asn
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " OD1",
        " ND2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HD2",
        "2HD2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 2  asn
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " OD1",
        " OD2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 3  asp
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " SG ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        " HG ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 4  cys
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD ",
        " OE1",
        " NE2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HG ",
        "2HG ",
        "1HE2",
        "2HE2",
        None,
        None,
        None,
        None,
        None,
    ),  # 5  gln
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD ",
        " OE1",
        " OE2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HG ",
        "2HG ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 6  glu
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        "1HA ",
        "2HA ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 7  gly
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " ND1",
        " CD2",
        " CE1",
        " NE2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "2HD ",
        "1HE ",
        "2HE ",
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 8  his
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG1",
        " CG2",
        " CD1",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        " HB ",
        "1HG2",
        "2HG2",
        "3HG2",
        "1HG1",
        "2HG1",
        "1HD1",
        "2HD1",
        "3HD1",
        None,
        None,
    ),  # 9  ile
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD1",
        " CD2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        " HG ",
        "1HD1",
        "2HD1",
        "3HD1",
        "1HD2",
        "2HD2",
        "3HD2",
        None,
        None,
    ),  # 10 leu
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD ",
        " CE ",
        " NZ ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HG ",
        "2HG ",
        "1HD ",
        "2HD ",
        "1HE ",
        "2HE ",
        "1HZ ",
        "2HZ ",
        "3HZ ",
    ),  # 11 lys
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " SD ",
        " CE ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HG ",
        "2HG ",
        "1HE ",
        "2HE ",
        "3HE ",
        None,
        None,
        None,
        None,
    ),  # 12 met
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD1",
        " CD2",
        " CE1",
        " CE2",
        " CZ ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HD ",
        "2HD ",
        "1HE ",
        "2HE ",
        " HZ ",
        None,
        None,
        None,
        None,
    ),  # 13 phe
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " HA ",
        "1HB ",
        "2HB ",
        "1HG ",
        "2HG ",
        "1HD ",
        "2HD ",
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 14 pro
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " OG ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HG ",
        " HA ",
        "1HB ",
        "2HB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 15 ser
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " OG1",
        " CG2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HG1",
        " HA ",
        " HB ",
        "1HG2",
        "2HG2",
        "3HG2",
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 16 thr
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD1",
        " CD2",
        " CE2",
        " CE3",
        " NE1",
        " CZ2",
        " CZ3",
        " CH2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HD ",
        "1HE ",
        " HZ2",
        " HH2",
        " HZ3",
        " HE3",
        None,
        None,
        None,
    ),  # 17 trp
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " CD1",
        " CD2",
        " CE1",
        " CE2",
        " CZ ",
        " OH ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "1HD ",
        "1HE ",
        "2HE ",
        "2HD ",
        " HH ",
        None,
        None,
        None,
        None,
    ),  # 18 tyr
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG1",
        " CG2",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        " HB ",
        "1HG1",
        "2HG1",
        "3HG1",
        "1HG2",
        "2HG2",
        "3HG2",
        None,
        None,
        None,
        None,
    ),  # 19 val
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "3HB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 20 unk
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "3HB ",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 21 mask
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " N9 ",
        " C4 ",
        " N3 ",
        " C2 ",
        " N1 ",
        " C6 ",
        " C5 ",
        " N7 ",
        " C8 ",
        " N6 ",
        None,
        None,
        "H5''",
        " H5'",
        " H4'",
        " H3'",
        "H2''",
        " H2'",
        " H1'",
        " H2 ",
        " H61",
        " H62",
        " H8 ",
        None,
        None,
    ),  # 22  DA
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " N1 ",
        " C2 ",
        " O2 ",
        " N3 ",
        " C4 ",
        " N4 ",
        " C5 ",
        " C6 ",
        None,
        None,
        None,
        None,
        "H5''",
        " H5'",
        " H4'",
        " H3'",
        "H2''",
        " H2'",
        " H1'",
        " H42",
        " H41",
        " H5 ",
        " H6 ",
        None,
        None,
    ),  # 23  DC
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " N9 ",
        " C4 ",
        " N3 ",
        " C2 ",
        " N1 ",
        " C6 ",
        " C5 ",
        " N7 ",
        " C8 ",
        " N2 ",
        " O6 ",
        None,
        "H5''",
        " H5'",
        " H4'",
        " H3'",
        "H2''",
        " H2'",
        " H1'",
        " H1 ",
        " H22",
        " H21",
        " H8 ",
        None,
        None,
    ),  # 24  DG
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " N1 ",
        " C2 ",
        " O2 ",
        " N3 ",
        " C4 ",
        " O4 ",
        " C5 ",
        " C7 ",
        " C6 ",
        None,
        None,
        None,
        "H5''",
        " H5'",
        " H4'",
        " H3'",
        "H2''",
        " H2'",
        " H1'",
        " H3 ",
        " H71",
        " H72",
        " H73",
        " H6 ",
        None,
    ),  # 25  DT
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        "H5''",
        " H5'",
        " H4'",
        " H3'",
        "H2''",
        " H2'",
        " H1'",
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 26  DX (unk DNA)
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " O2'",
        " N1 ",
        " C2 ",
        " N3 ",
        " C4 ",
        " C5 ",
        " C6 ",
        " N6 ",
        " N7 ",
        " C8 ",
        " N9 ",
        None,
        " H5'",
        "H5''",
        " H4'",
        " H3'",
        " H2'",
        "HO2'",
        " H1'",
        " H2 ",
        " H61",
        " H62",
        " H8 ",
        None,
        None,
    ),  # 27   A
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " O2'",
        " N1 ",
        " C2 ",
        " O2 ",
        " N3 ",
        " C4 ",
        " N4 ",
        " C5 ",
        " C6 ",
        None,
        None,
        None,
        " H5'",
        "H5''",
        " H4'",
        " H3'",
        " H2'",
        "HO2'",
        " H1'",
        " H42",
        " H41",
        " H5 ",
        " H6 ",
        None,
        None,
    ),  # 28   C
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " O2'",
        " N1 ",
        " C2 ",
        " N2 ",
        " N3 ",
        " C4 ",
        " C5 ",
        " C6 ",
        " O6 ",
        " N7 ",
        " C8 ",
        " N9 ",
        " H5'",
        "H5''",
        " H4'",
        " H3'",
        " H2'",
        "HO2'",
        " H1'",
        " H1 ",
        " H22",
        " H21",
        " H8 ",
        None,
        None,
    ),  # 29   G
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " O2'",
        " N1 ",
        " C2 ",
        " O2 ",
        " N3 ",
        " C4 ",
        " O4 ",
        " C5 ",
        " C6 ",
        None,
        None,
        None,
        " H5'",
        "H5''",
        " H4'",
        " H3'",
        " H2'",
        "HO2'",
        " H1'",
        " H3 ",
        " H5 ",
        " H6 ",
        None,
        None,
        None,
    ),  # 30   U
    (
        " O4'",
        " C1'",
        " C2'",
        " OP1",
        " P  ",
        " OP2",
        " O5'",
        " C5'",
        " C4'",
        " C3'",
        " O3'",
        " O2'",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H5'",
        "H5''",
        " H4'",
        " H3'",
        " H2'",
        "HO2'",
        " H1'",
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # 31  RX (unk RNA)
    (
        " N  ",
        " CA ",
        " C  ",
        " O  ",
        " CB ",
        " CG ",
        " NE2",
        " CD2",
        " CE1",
        " ND1",
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        " H  ",
        " HA ",
        "1HB ",
        "2HB ",
        "2HD ",
        "1HE ",
        "1HD ",
        None,
        None,
        None,
        None,
        None,
        None,
    ),  # -1 his_d
]

# frames for generic FAPE
##NOTE: 1st entry must be "backbone frame"
# fmt: on
frames = [
    [[" N  ", " CA ", " C  "], [" CA ", " C  ", " O  "]],  # ala
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
        [" CB ", " CG ", " CD "],
        [" CG ", " CD ", " NE "],
    ],  # arg
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
    ],  # asn
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
    ],  # asp
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
    ],  # cys
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
        [" CB ", " CG ", " CD "],
    ],  # gln
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
        [" CB ", " CG ", " CD "],
    ],  # glu
    [[" N  ", " CA ", " C  "], [" CA ", " C  ", " O  "]],  # gly
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
    ],  # his
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG1"],
    ],  # ile
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
    ],  # leu
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
        [" CB ", " CG ", " CD "],
        [" CG ", " CD ", " CE "],
    ],  # lys
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
        [" CB ", " CG ", " SD "],
    ],  # met
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
    ],  # phe
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
        [" CB ", " CG ", " CD "],
    ],  # pro
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " OG "],
    ],  # ser
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " OG1"],
    ],  # thr
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
    ],  # trp
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
        [" CA ", " CB ", " CG "],
    ],  # tyr
    [
        [" N  ", " CA ", " C  "],
        [" CA ", " C  ", " O  "],
        [" N  ", " CA ", " CB "],
    ],  # val
    [[" N  ", " CA ", " C  "], [" CA ", " C  ", " O  "]],  # unk
    [[" N  ", " CA ", " C  "], [" CA ", " C  ", " O  "]],  # mask
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N9 "],
        [" C4'", " C3'", " O3'"],
    ],  # DA
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N1 "],
        [" C4'", " C3'", " O3'"],
    ],  # DC
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N9 "],
        [" C4'", " C3'", " O3'"],
    ],  # DG
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N1 "],
        [" C4'", " C3'", " O3'"],
    ],  # DT
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C4'", " C3'", " O3'"],
    ],  # DX
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N9 "],
        [" C4'", " C3'", " O3'"],
    ],  # A
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N1 "],
        [" C4'", " C3'", " O3'"],
    ],  # C
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N9 "],
        [" C4'", " C3'", " O3'"],
    ],  # G
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C2'", " C1'", " N1 "],
        [" C4'", " C3'", " O3'"],
    ],  # U
    [
        [" O4'", " C1'", " C2'"],
        [" OP1", " P  ", " O5'"],
        [" P  ", " O5'", " C5'"],
        [" O5'", " C5'", " C4'"],
        [" C5'", " C4'", " C3'"],
        [" C5'", " C4'", " O4'"],
        [" C4'", " O4'", " C1'"],
        [" C4'", " C3'", " O3'"],
    ],  # RX
]
# fmt: on
NFRAMES = max([len(f) for f in frames])

# general FAPE parameters
frame_indices = torch.full((NAATOKENS, NFRAMES, 3, 2), 0, dtype=torch.long)
for i in range(NNAPROTAAS):
    i_l = aa2long[i]
    for j, x in enumerate(frames[i]):
        if x is not None:
            # frames are stored as (residue offset, atom position)
            frame_indices[i, j, 0] = torch.tensor((0, i_l.index(x[0])))
            frame_indices[i, j, 1] = torch.tensor((0, i_l.index(x[1])))
            frame_indices[i, j, 2] = torch.tensor((0, i_l.index(x[2])))


allatom_mask = torch.zeros((NAATOKENS, NTOTAL), dtype=torch.bool)
long2alt = torch.zeros((NAATOKENS, NTOTAL), dtype=torch.long)
for i in range(NNAPROTAAS):
    for j, a in enumerate(aa2long[i]):
        if a is not None:
            allatom_mask[i, j] = True

allatom_mask[NNAPROTAAS:, 1] = True
allatom_mask[:, 1] = True
heavyatom_mask = allatom_mask.clone()
heavyatom_mask[:, NHEAVY:] = False
